
def get_repetition_penalty_reward(ngram_size: int = 3, max_penalty: float = -0.1):
    if max_penalty > 0:
        raise ValueError(f"max_penalty {max_penalty} should not be positive")
    def zipngram(text: str, ngram_size: int):
        words = text.lower().split()
        return zip(*[words[i:] for i in range(ngram_size)])
    def repetition_penalty_reward(completions,**kwargs) -> float:
        contents = [completion[0]["content"] for completion in completions]
        rewards =[]
        for completion in contents:
            if completion == "":
                rewards.append(0.0)
                continue
            if len(completion.split()) < ngram_size:
                rewards.append(0.0)
                continue
            ngrams = set()
            total = 0
            for ng in zipngram(completion,ngram_size):
                ngrams.add(ng)
                total +=1
            scaling = 1-len(ngrams) / total
            reward = scaling * max_penalty
            rewards.append(reward)
        return rewards
    return get_repetition_penalty_reward